#!/usr/bin/env python

import os, sys
import glob
from astropy.visualization import ZScaleInterval
import numpy as np
import matplotlib.pylab as plt
from astropy.io import fits

wave = [300.0, 350.0, 390.0, 450.0, 490.0, 530.0, 570.0, 610.0, 650.0, 650.0, 710]

print(wave)

fitslist = glob.glob(sys.argv[1]+'/*.fits')

flist = []


for i in range(len(wave)):
    for j in range(len(fitslist)):
        if fitslist[j].find('{}nm'.format(wave[i])) >= 0:
            flist.append(fitslist[j])

print(flist)
#sys.exit(0)

flip = False


f_ref = fits.getdata(flist[9])

mask1 = f_ref > np.median(f_ref) + np.std(f_ref)*5
mask2 = f_ref < np.median(f_ref) - np.std(f_ref)*5
f_ref[mask1] = np.median(f_ref)
f_ref[mask2] = np.median(f_ref)

nx, ny = f_ref.shape[0], f_ref.shape[1]
if flip:
    f_ref_os4 = f_ref[(nx//2):, 1152*3:1152*4]
    f_ref_os4 = np.flip(f_ref_os4,axis=0)
    f_ref[0:nx//2, 1152*3:1152*4] = f_ref_os4

fig, ax = plt.subplots(2, 3, sharex=True, sharey=True,
                       constrained_layout=True, figsize=(8,5))
ax = ax.flatten()

zs = ZScaleInterval()
for i in range(6):
    print('--> loading {}'.format(flist[i]))
    f = fits.getdata(flist[i])

    mask1 = f > np.median(f) + np.std(f)*5
    mask2 = f < np.median(f) - np.std(f)*5
    f[mask1] = np.median(f)
    f[mask2] = np.median(f)
    if flip:
        fos4 = f[(nx//2):, 1152*3:1152*4]
        fos4 = np.flip(fos4,axis=0)
        f[0:nx//2, 1152*3:1152*4] = fos4

    # for j in range(8):
    #     f[0:4616,j*1152:(j+1)*1152] /= np.median(f[0:4616,j*1152:(j+1)*1152])
    #     f[4616:,j*1152:(j+1)*1152] /= np.median(f[4616:,j*1152:(j+1)*1152])

    # img = f/np.median(f)
    img = f/f_ref
    img = img / np.median(img)
    index = np.isfinite(img)
    vmin, vmax = zs.get_limits(img[index])
    im = ax[i].imshow(img, vmin=vmin, vmax=vmax, cmap='Blues_r', origin='lower')
    ax[i].set_title('{}nm'.format(wave[i]))
    ax[i].set_xticklabels([])
    ax[i].set_yticklabels([])

fig.colorbar(im, ax=[ax[i] for i in range(6)], fraction=0.025, pad=0.05, aspect=30)
plt.savefig('test.png', bbox_inches='tight')
plt.show()
